#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <limits.h>
#include <sys/stat.h>

extern char **environ;

static int file_exists(const char *path)
{
    struct stat st;
    return (stat(path, &st) == 0);
}

static void build_preload_path(char *out, size_t outsz)
{
    ssize_t n = readlink("/proc/self/exe", out, outsz - 1);
    if (n <= 0) { out[0] = '\0'; return; }
    out[n] = '\0';
    char *last = strrchr(out, '/');
    if (!last) { out[0] = '\0'; return; }
    *(last) = '\0';
    size_t len = strlen(out);
    snprintf(out + len, outsz - len, "/libmemcheck_preload.z.so");
}

static char **make_env_with_preload(const char *preload_so)
{
    size_t cnt = 0; for (char **p = environ; p && *p; ++p) cnt++;
    char **env = (char **)calloc(cnt + 2, sizeof(char *));
    if (!env) return NULL;
    size_t j = 0; int replaced = 0;
    for (size_t i = 0; i < cnt; ++i) {
        if (strncmp(environ[i], "LD_PRELOAD=", 11) == 0) {
            size_t need = strlen("LD_PRELOAD=") + strlen(preload_so) + 1;
            char *v = (char *)malloc(need);
            if (!v) { free(env); return NULL; }
            snprintf(v, need, "LD_PRELOAD=%s", preload_so);
            env[j++] = v; replaced = 1;
        } else {
            env[j++] = environ[i];
        }
    }
    if (!replaced) {
        size_t need = strlen("LD_PRELOAD=") + strlen(preload_so) + 1;
        char *v = (char *)malloc(need);
        if (!v) { free(env); return NULL; }
        snprintf(v, need, "LD_PRELOAD=%s", preload_so);
        env[j++] = v;
    }
    env[j] = NULL;
    return env;
}

int main(int argc, char **argv)
{
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <program> [args...]\n", argv[0]);
        return 2;
    }
    char so_path[PATH_MAX] = {0};
    build_preload_path(so_path, sizeof(so_path));
    const char *preload_so = so_path[0] ? so_path : "libmemcheck_preload.z.so";
    if (!file_exists(preload_so)) {
        fprintf(stderr, "memcheck: preload so not found: %s\n", preload_so);
        return 1;
    }
    char **new_env = make_env_with_preload(preload_so);
    if (!new_env) {
        fprintf(stderr, "memcheck: failed to set LD_PRELOAD\n");
        return 1;
    }
    char **child_argv = &argv[1];
    execve(child_argv[0], child_argv, new_env);
    perror("execve");
    return 127;
}


